#!/usr/bin/env python3
'''
未验证
'''
import gymnasium as gym
import ptan
import math
import numpy as np
import time
import argparse
import collections
from tensorboardX import SummaryWriter

import torch
import torch.nn.utils as nn_utils
import torch.nn.functional as F
import torch.optim as optim
import torch.multiprocessing as mp
import os

from lib import common, model

GAMMA = 0.99
LEARNING_RATE = 0.001
ENTROPY_BETA = 0.01
# 训练数据样本缓冲区的长度,这里之所以训练数据缓冲区的长度比之前的长
# 是因为多个环境并行获取数据，缓冲区内的训练数据基本都是最新的，如果是单个
# 环境下，这么长的缓冲区，其必定会有过旧的训练数据
BATCH_SIZE = 128

REWARD_STEPS = 2
CLIP_GRAD = 0.1

# 创建子进程的总数
PROCESSES_COUNT = 4
NUM_ENVS = 15

def calc_logprob(mu_v, var_v, actions_v):
    '''
    实现P292页的公式
    用来实现计算执行动作的概率密度
    所谓的概率密度，指的是在高斯函数中，在x点附近取值的概率值，高斯概率密度函数计算的仅仅只是一个点的概率值，要计算实际的概率大小需要结合实际的面积才能得到，而这里不需要计算一整个概率大小，而是计算概率密度即可，用概率密度大小替代实际的概率也是可以的
    根据高斯分布的概率密度函数，所选择的动作离高斯分布的均值越远，则概率越小，选择的动作离高斯分布的均值越近，则概率越大；
    对于方差来说，方差越小，则概率越大；方差越大，则概率越小；
    所以如果要使得计算的log_prob_v最大话，对于正向的adv_v，说明当前执行的动作具备优势，就会将均值靠近当前执行的动作，方差也会缩小；对于负向的adv_v，说明当前执行的动作具备劣势，就会将均值远离当前执行的动作，方差也会放大；
    这样计算更新下去，最后使得计算的均值离执行的正确动作越来越近，方差也越来越小
    '''
    p1 = - ((mu_v - actions_v) ** 2) / (2*var_v.clamp(min=1e-3))
    p2 = - torch.log(torch.sqrt(2 * math.pi * var_v))
    return p1 + p2

def make_env():
    # 创建环境，将不同环境统一成统一操作
    return gym.make("Swimmer-v4")

# 创建一个可以给成员命名的元组，可以类似于类一样使用
TotalReward = collections.namedtuple('TotalReward', field_names='reward')


def data_func(net, device, train_queue):
    '''
    子进程主要执行的动作是
    不断的从环境中获取数据，并将其存储到共享游戏数据队列中
    '''

    # 每个子进程创建NUM_ENVS个环境
    envs = [make_env() for _ in range(NUM_ENVS)]
    # 创建策略代理
    agent = model.AgentA2C(net, device=device)
    # 创建经验源，经验源本身支持多环境
    exp_source = ptan.experience.ExperienceSourceFirstLast(envs, agent, gamma=GAMMA, steps_count=REWARD_STEPS)

    for exp in exp_source:
        # 获取经验源中所有的回报
        # pop_total_rewards存在返回值，那么表示一次游戏
        # 已经结束
        new_rewards = exp_source.pop_total_rewards()
        if new_rewards:
            # 将游戏结束奖励存储到train_queue中
            # np.mean计算游戏统计游戏多次结束的奖励平均值
            train_queue.put(TotalReward(reward=np.mean(new_rewards)))
        # 如果游戏没有结束，则将目前进行的游戏记录的数据记录在train_queue中
        train_queue.put(exp)


def test_net(net, env, device, count=10):
    """测试网络在环境中的表现，返回平均的奖励值"""
    rewards = 0.0
    steps = 0
    for _ in range(count):
        obs = env.reset()
        while True:
            obs_v = torch.tensor([obs], dtype=torch.float32).to(device)
            mu_v, _, _ = net(obs_v)
            action = mu_v.squeeze(dim=0).data.cpu().numpy()
            obs, reward, terminated, truncated, _ = env.step(action)
            rewards += reward
            steps += 1
            if terminated or truncated:
                break
    return rewards / count, steps // count


if __name__ == "__main__":
    # 该方法主要有两个作用:
    #
    # 设置多进程的启动方式(spawn或fork)。默认是spawn。
    # 避免多进程模块的多次导入导致的错误。
    # 关于启动方式:
    #
    # spawn: 子进程将只继承父进程的信号处理机制,而不会继承父进程的句柄、锁等资源。所以spawn方式更安全,也是推荐的默认方式。
    # fork: 子进程通过fork启动,会通过复制父进程获得父进程的全部资源,包括锁、信号处理等,所以可能会导致在多进程中出现问题。
    # 由于pytorch的限制，spawn是最佳的方式
    mp.set_start_method('spawn')
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=True, action="store_true", help="Enable cuda")
    parser.add_argument("-n", "--name", default='a3c-data', help="Name of the run")
    args = parser.parse_args()
    device = "cuda" if args.cuda else "cpu"
    save_path = os.path.join("saves", "a3c-" + args.name)
    os.makedirs(save_path, exist_ok=True)

    # 为记录统计设置名字
    writer = SummaryWriter(comment="-a3c-data_" + args.name)

    env = make_env()
    test_env = make_env()
    net = model.ModelA2C(env.observation_space.shape[0], env.action_space.shape[0]).to(device)
    # 生命网络net为共享内存，share_memory必须在调用fork前调用
    # 这样net就可以很方便的在每个进程间使用了
    # 如果net设置的dev是cuda，那么网络张量默认的共享的，不用调用share_memory
    # 如果net设置的dev是cpu，那么网络张量是不共享的，需要调用share_memory
    net.share_memory()

    optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE, eps=1e-3)

    # 创建子进程共享队列，最大存储的数量为子进程的数量
    train_queue = mp.Queue(maxsize=PROCESSES_COUNT)
    data_proc_list = []
    for _ in range(PROCESSES_COUNT):
        # 创建子进程，将net、目标设备、训练子进程队列传入子进程的参数中
        data_proc = mp.Process(target=data_func, args=(net, device, train_queue))
        data_proc.start()
        data_proc_list.append(data_proc)

    batch = []
    best_reward = None  # 记录最佳回报
    # 检查是否存在已保存的模型
    step_idx = 0
    if os.path.exists(save_path) and len(os.listdir(save_path)) > 0:
        checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(save_path)),
                             key=lambda x: int(x.split('_')[1].split('.')[0]))
        checkpoint = torch.load(os.path.join(save_path, checkpoints[-1]), map_location=device, weights_only=False)
        net.load_state_dict(checkpoint['net'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        step_idx = checkpoint['step_idx']
        print(f"成功加载模型: {checkpoints[-1]}")

    try:
        # todo 完善RewardTracker注释 作用统计近100次游戏的奖励，如果奖励达标则退出游戏
        with common.RewardTracker(writer, stop_reward=10000) as tracker:
            # todo 完善TBMeanTracker注释，作用将最近的10步游戏记录到tensorboard中
            with ptan.common.utils.TBMeanTracker(writer, batch_size=100) as tb_tracker:
                while True:
                    # 从子进程获取环境数据
                    train_entry = train_queue.get()
                    # todo 从代码可知，train_entry有两种类型，其中一种类型是TotalReward
                    #  TotalReward内部只有奖励回报，为什么可以混在一起不会出错
                    if isinstance(train_entry, TotalReward):
                        # 如果感知到有一次游戏结束，则判断获取的奖励是否符合训练结束的条件
                        if tracker.reward(train_entry.reward, step_idx):
                            break
                        continue

                    step_idx += 1
                    batch.append(train_entry)
                    if len(batch) < BATCH_SIZE:
                        continue

                    states_v, actions_v, vals_ref_v = \
                        common.unpack_batch_a2c(batch, net, last_val_gamma=GAMMA**REWARD_STEPS, device=device)
                    batch.clear()

                    # 清空梯度
                    optimizer.zero_grad()
                    # 得到动作预测，当前状态下的Q值
                    mu_v, var_v, value_v = net(states_v)

                    # 训练网络预测Q值，预测出来的值和bellman计算出来的Q值应该接近
                    loss_value_v = F.mse_loss(value_v.squeeze(-1), vals_ref_v)

                    adv_v = vals_ref_v.unsqueeze(dim=-1) - value_v.detach()
                    # 这里就对应之前策略梯度中计算每个动作的奖励以及对应动作的概率的损失
                    # 为什么采用高斯分布，而不是直接网络输出预测目标值val，然后目标值与执行的动作值做平方差等方式计算
                    # 1. 计算连续值的概率的公式之一就是高斯分布
                    # 2. 高斯概率计算公式输出均值和方差，均值可以代表目标动作，方差代表预测动作是否稳定，通过这两个值既可以实现探索，也可以促进网络往稳定的方向训练（但是这里并没有使用方差去进行探索），而是采用了信息熵的方式去促使网络进行探索
                    log_prob_v = adv_v * calc_logprob(mu_v, var_v, actions_v)
                    # 去负值，计算最小值
                    loss_policy_v = -log_prob_v.mean()
                    # 高斯概率密度分度熵的计算公式 P292页，同样这里也取了负值
                    entropy_loss_v = ENTROPY_BETA * (-(torch.log(2*math.pi*var_v) + 1)/2).mean()

                    # 损失之和，计算最小值，更细梯度梯度
                    loss_v = loss_policy_v + entropy_loss_v + loss_value_v
                    loss_v.backward()
                    optimizer.step()

                    # 将相关的值记录起来
                    tb_tracker.track("advantage", adv_v, step_idx)
                    tb_tracker.track("values", value_v, step_idx)
                    tb_tracker.track("batch_rewards", vals_ref_v, step_idx)
                    tb_tracker.track("loss_entropy", entropy_loss_v, step_idx)
                    tb_tracker.track("loss_policy", loss_policy_v, step_idx)
                    tb_tracker.track("loss_value", loss_value_v, step_idx)
                    tb_tracker.track("loss_total", loss_v, step_idx)

                    # 每隔一定步数保存一次模型
                    if step_idx % 1000 == 0:
                        ts = time.time()
                        rewards, steps = test_net(net, test_env, device=device)
                        print(f"Test done in {time.time() - ts:.2f} sec, reward {rewards:.3f}, steps {steps}")
                        writer.add_scalar("test_reward", rewards, step_idx)
                        writer.add_scalar("test_steps", steps, step_idx)

                        # 保存最好的模型
                        if best_reward is None or best_reward < rewards:
                            if best_reward is not None:
                                print(f"Best reward updated: {best_reward:.3f} -> {rewards:.3f}")
                                checkpoint_name = os.path.join(save_path, f"best_{rewards:+.3f}_{step_idx}.pth")
                                torch.save(net.state_dict(), checkpoint_name)
                            best_reward = rewards

                        checkpoint = {
                            'net': net.state_dict(),
                            'optimizer': optimizer.state_dict(),
                            'step_idx': step_idx
                        }
                        checkpoint_name = os.path.join(save_path, f'epoch_{step_idx}.pth')
                        common.save_checkpoints(step_idx, checkpoint, save_path, "a3c", keep_last=5)
                        print(f"保存模型: {checkpoint_name}")
    finally:
        # 防止出现异常时，没有管子子进程
        for p in data_proc_list:
            p.terminate()
            p.join()
